package com.hades.hash.consistent;

import java.net.InetAddress;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

/**
 * @program: hash
 * @description:
 * @author: hadesinb
 * @create: 2021-05-20 10:21
 */
public class ConsistentHash {

    private final TreeMap<Long, InetAddress> virtualInvokers;

    public ConsistentHash(List<InetAddress> invokers) {
        this.virtualInvokers = new TreeMap<>();
        int replicaNumber = 180;
        for (InetAddress address : invokers) {
            String addr = address.getHostAddress();
            for (int i = 0; i < replicaNumber; i++) {
                byte[] digest = md5(addr + i);
                for (int h = 0; h < 4; h++) {
                    long m = hash(digest, h);
                    virtualInvokers.put(m, address);
                }
            }
        }
    }

    public InetAddress select(String key) {
        byte[] digest = md5(key);
        return selectForKey(hash(digest, 0));
    }

    private InetAddress selectForKey(long hash) {
        Map.Entry<Long, InetAddress> entry = virtualInvokers.ceilingEntry(hash);
        if (entry == null) {
            entry = virtualInvokers.firstEntry();
        }
        return entry.getValue();
    }

    private long hash(byte[] digest, int number) {
        return (((long) (digest[3 + number * 4] & 0xFF) << 24)
                | ((long) (digest[2 + number * 4] & 0xFF) << 16)
                | ((long) (digest[1 + number * 4] & 0xFF) << 8)
                | (digest[number * 4] & 0xFF))
                & 0xFFFFFFFFL;
    }

    private byte[] md5(String value) {
        MessageDigest md5;
        try {
            md5 = MessageDigest.getInstance("MD5");
        } catch (NoSuchAlgorithmException e) {
            throw new IllegalStateException(e.getMessage(), e);
        }
        md5.reset();
        byte[] bytes = value.getBytes(StandardCharsets.UTF_8);
        md5.update(bytes);
        return md5.digest();
    }
}
